import datetime
import time
import json
from typing import Dict, Callable
from ..models import WSSpeedControlMsg, WSMessageTypes
from flask_sock import Sock, Server, ConnectionClosed
from element_sci_api.api import BaseFlaskSock
import logging

logger = logging.getLogger(__name__)


class ControlWebsocket(BaseFlaskSock):
    def __init__(self, sock, rel_url):
        super().__init__(sock, rel_url)
        self._on_msg_handlers: Dict[WSMessageTypes,
                                    Callable[[ControlWebsocket, dict], None]] = {}

    def register_msg_handler(self, msg_type: WSMessageTypes, msg_handler:  Callable[["ControlWebsocket", dict], None]):
        """
        注册消息处理函数
        """
        self._on_msg_handlers[msg_type] = msg_handler

    def on_open(self, ws):
        self.send_data("bbbbaaaaaaa")

    def on_message(self, msg: bytes):
        try:
            data = json.loads(msg)
            if 'type' in data:
                self._on_msg_handlers[data['type']](self, data)
            else:
                logger.error(f"Cannot parse this message: {data}")
        except json.JSONDecodeError:
            logger.error(f"Cannot parse this message: {data}")
